#include "toroidal.h"
#include "stdlib.h"
#include <stdio.h>

field::field()
{
    for(int i = 0; i<Nx*Ny; i++)
        this->data[i] = 0.0;
}

vfield::vfield()
{
    this->fr = new double[Nx*Ny];
    this->fz = new double[Nx*Ny];
    this->fphi = new double[Nx*Ny];
    for(int i = 0; i<Nx*Ny; i++)
    {
        this->fr[i] = 0.0;
        this->fz[i] = 0.0;
        this->fphi[i] = 0.0;
    }
}

field::field(double* dt) //copy data from pointer to array
{
    for(int i = 0; i<Nx*Ny; i++)
        this->data[i] = *(dt+i);
}

void partialy(field a, double* parx)
{
    double *scaf = a.getdata();
    for(int i = 0; i < Nx*Ny; i++)
    {
        int ix = i/Ny;
        int iy = i%Ny;
        if(iy ==0)
            parx[i] = (scaf[i+1] - scaf[i])/dz;
        else if (iy == Ny-1)
            parx[i] = (scaf[i] - scaf[i-1])/dz;
        else
            parx[i] = (scaf[i+1] - scaf[i-1])/(2*dz);
    }
}

void partialx(field a, double* pary)
{
    double *scaf = a.getdata();
    for(int i = 0; i < Nx*Ny; i++)
    {
        int ix = i/Ny;
        if(ix ==0)
            pary[i] = (scaf[i+Ny] - scaf[i])/dr;
        else if (ix == Nx-1)
            pary[i] = (scaf[i] - scaf[i-Ny])/dr;
        else
            pary[i] = (scaf[i+Ny] - scaf[i-Ny])/(2*dr);
    }
}

vfield field::gradient()
{
    double *parx = new double[Nx*Ny];
    double *pary = new double[Nx*Ny];
    double *parphi = new double[Nx*Ny];

    for(int i=0; i <Nx*Ny; i++)
        parphi[i] = 0.0;
    partialx(*this, parx);
    partialy(*this, pary);
    
    field parr = field(parx);
    field parz = field(pary);
    field pphi = field(parphi);
    vfield vf = vfield(&parr,&pphi,&parz);
    return vf;
}   

void field::scale(double fact)
{
    for(int i = 0; i < Nx*Ny; i++)
    {
        this->data[i] *= fact;
    }
}
vfield::vfield(field* a, field* b, field* c)
{
    this->fr = new double[Nx*Ny];
    this->fz = new double[Nx*Ny];
    this->fphi = new double[Nx*Ny];
    double* pr = a->getdata();
    double* pz = c->getdata();
    double* pphi = b->getdata();
    for(int i = 0; i < Nx*Ny; i++)
    {
        this->fr[i] = *(pr+i);
        this->fz[i] = *(pz+i);
        this->fphi[i] = *(pphi+i);
    }
}

vfield::vfield(double* a, double* b, double* c)
{
    this->fr = new double[Nx*Ny];
    this->fz = new double[Nx*Ny];
    this->fphi = new double[Nx*Ny];
    for(int i = 0; i < Nx*Ny; i++)
    {
        this->fr[i] = *(a+i);
        this->fphi[i] = *(b+i);
        this->fz[i] = *(c+i);
    }
}

vec vfield::interp(cylv cord)
{   
    vec ret = {0.0, 0.0, 0.0};
    double r = cord.r;
    double z = cord.z;
    int rg = (int)((r - Rmin)/dr);
    int zg = (int)((z - Zmin)/dz);
    double ofr = (r - (double)rg * dr - Rmin)/dr;
    double ofz = (z - (double)zg * dz - Zmin)/dz;
    double coefrs[4] = {ofr*(ofr-1.0)*(ofr-2.0)/-6.0, (ofr+1.0)*(ofr-1.0)*(ofr-2.0)/2.0,
    (ofr+1.0)*ofr*(ofr-2.0)/-2.0, (ofr+1.0)*ofr*(ofr-1.0)/6.0};
    double coefzs[4] = {ofz*(ofz-1.0)*(ofz-2.0)/-6.0, (ofz+1.0)*(ofz-1.0)*(ofz-2.0)/2.0,
    (ofz+1.0)*ofz*(ofz-2.0)/-2.0, (ofz+1.0)*ofz*(ofz-1.0)/6.0};
    for(int i = 0; i < 4; i++)
        for(int j = 0; j < 4; j++)
        {
            ret.x += coefrs[i]*coefzs[j]*fr[Ny*(rg+i-1)+zg+j-1];
            ret.y += coefrs[i]*coefzs[j]*fphi[Ny*(rg+i-1)+zg+j-1];
            ret.z += coefrs[i]*coefzs[j]*fz[Ny*(rg+i-1)+zg+j-1];
        }
    // rotation needs to be considered
    return ret;
}

void vfield::curl(vfield* vfp) // write to vfp pointer
{
    double *curlR = vfp->getfr();
    double *curlZ = vfp->getfz();
    double *curlphi = vfp->getfphi();

    double *ar = this->fr;
    double *az = this->fz;
    double *ap = this->fphi;

    double* prfz = new double[Nx*Ny];
    double* pzfr = new double[Nx*Ny];
    double* pzfp = new double[Nx*Ny];
    double* prfp = new double[Nx*Ny];

    partialx(field(az),prfz);
    partialy(field(ar),pzfr);
    partialy(field(ap),pzfp);
    partialx(field(ap),prfp);

    for (int i = 0; i<Nx*Ny; i++)
    {
        double R = (i/Ny)*dr + Rmin;
        curlphi[i] = -prfz[i]+pzfr[i];
        curlR[i] = -pzfp[i];
        curlZ[i] = prfp[i] + ap[i]/R;
    }
}

void vfield::scale(double x)
{
    for(int i = 0; i < Nx*Ny; i++)
    {
        this->fr[i] *= x;
        this->fz[i] *= x;
        this->fphi[i] *= x;
    }
}

field fadd(field a, field b)
{
    double* added = new double[Nx*Ny];
    double* sca = a.getdata();
    double* scb = b.getdata();
    for(int i = 0; i < Nx*Ny; i++)
    {
        added[i] = sca[i] + scb[i];
    }
    return field(added);
}

field fmul(field a, field b)
{
    double* muled = new double[Nx*Ny];
    double* sca = a.getdata();
    double* scb = b.getdata();
    for(int i = 0; i < Nx*Ny; i++)
    {
        muled[i] = sca[i]*scb[i];
    }
    return field(muled);
}

